# Author: WY
# Date: 2022/11/10 11:09

from sklearn import preprocessing
from torch.utils.data import Dataset, DataLoader
import numpy as np

# 数据预处理MinMaxScaler(数据范围缩放到[-1, 1]之间)
minmaxScaler = preprocessing.MinMaxScaler(feature_range=(0., 1.))


# Dataset
class my_Dataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data[0])

    def __getitem__(self, item):
        return self.data[:, item, :]


def get_loader(split: float, batch_size: int, data):

    data = np.array(data)

    # 归一化特征数据
    data = np.array([minmaxScaler.fit_transform(item) for item in data])

    data_len = len(data[0])
    split_line = int(data_len * split)

    # DataSet和DataLoader
    dataset_train = my_Dataset(data[:, 0: split_line, :])
    dataset_test = my_Dataset(data[:, split_line:, :])
    dataloader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=False, drop_last=True)
    dataloader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False, drop_last=True)
    return dataloader_train, dataloader_test
